{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "  import pydrake\n",
    "  import underactuated\n",
    "except ImportError:\n",
    "  !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "  from jupyter_setup import setup_underactuated\n",
    "  setup_underactuated()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "from pydrake.solvers.mathematicalprogram import MathematicalProgram, Solve\n",
    "from pydrake.trajectories import PiecewisePolynomial\n",
    "\n",
    "from underactuated.quadrotor2d import Quadrotor2DVisualizer\n",
    "from underactuated.jupyter import SetupMatplotlibBackend\n",
    "plt_is_interactive = SetupMatplotlibBackend()\n",
    "\n",
    "# TODO(russt): Use drake.trajectories.PiecewisePolynomialTrajectory\n",
    "#  instead (currently missing python bindings for the required constructor),\n",
    "#  or port this class to C++.\n",
    "class PPTrajectory():\n",
    "\n",
    "    def __init__(self, sample_times, num_vars, degree, continuity_degree):\n",
    "        self.sample_times = sample_times\n",
    "        self.n = num_vars\n",
    "        self.degree = degree\n",
    "\n",
    "        self.prog = MathematicalProgram()\n",
    "        self.coeffs = []\n",
    "        for i in range(len(sample_times)):\n",
    "            self.coeffs.append(\n",
    "                self.prog.NewContinuousVariables(num_vars, degree + 1, \"C\"))\n",
    "        self.result = None\n",
    "\n",
    "        # Add continuity constraints\n",
    "        for s in range(len(sample_times) - 1):\n",
    "            trel = sample_times[s + 1] - sample_times[s]\n",
    "            coeffs = self.coeffs[s]\n",
    "            for var in range(self.n):\n",
    "                for deg in range(continuity_degree + 1):\n",
    "                    # Don't use eval here, because I want left and right\n",
    "                    # values of the same time\n",
    "                    left_val = 0\n",
    "                    for d in range(deg, self.degree + 1):\n",
    "                        left_val += (coeffs[var, d] * np.power(trel, d - deg) *\n",
    "                                     math.factorial(d) /\n",
    "                                     math.factorial(d - deg))\n",
    "                    right_val = (self.coeffs[s + 1][var, deg] *\n",
    "                                 math.factorial(deg))\n",
    "                    self.prog.AddLinearConstraint(left_val == right_val)\n",
    "\n",
    "        # Add cost to minimize highest order terms\n",
    "        for s in range(len(sample_times) - 1):\n",
    "            self.prog.AddQuadraticCost(np.eye(num_vars), np.zeros(\n",
    "                (num_vars, 1)), self.coeffs[s][:, -1])\n",
    "\n",
    "    def eval(self, t, derivative_order=0):\n",
    "        if derivative_order > self.degree:\n",
    "            return 0\n",
    "\n",
    "        s = 0\n",
    "        while s < len(self.sample_times) - 1 and t >= self.sample_times[s + 1]:\n",
    "            s += 1\n",
    "        trel = t - self.sample_times[s]\n",
    "\n",
    "        if self.result is None:\n",
    "            coeffs = self.coeffs[s]\n",
    "        else:\n",
    "            coeffs = self.result.GetSolution(self.coeffs[s])\n",
    "\n",
    "        deg = derivative_order\n",
    "        val = 0 * coeffs[:, 0]\n",
    "        for var in range(self.n):\n",
    "            for d in range(deg, self.degree + 1):\n",
    "                val[var] += (coeffs[var, d] * np.power(trel, d - deg) *\n",
    "                             math.factorial(d) / math.factorial(d - deg))\n",
    "\n",
    "        return val\n",
    "\n",
    "    def add_constraint(self, t, derivative_order, lb, ub=None):\n",
    "        \"\"\"Adds a constraint of the form d^deg lb <= x(t) / dt^deg <= ub.\"\"\"\n",
    "        if ub is None:\n",
    "            ub = lb\n",
    "\n",
    "        assert derivative_order <= self.degree\n",
    "        val = self.eval(t, derivative_order)\n",
    "        self.prog.AddLinearConstraint(val, lb, ub)\n",
    "\n",
    "    def generate(self):\n",
    "        self.result = Solve(self.prog)\n",
    "        assert self.result.is_success()\n",
    "\n",
    "\n",
    "tf = 3\n",
    "zpp = PPTrajectory(sample_times=np.linspace(0, tf, 6),\n",
    "                   num_vars=2,\n",
    "                   degree=5,\n",
    "                   continuity_degree=4)\n",
    "zpp.add_constraint(t=0, derivative_order=0, lb=[0, 0])\n",
    "zpp.add_constraint(t=0, derivative_order=1, lb=[0, 0])\n",
    "zpp.add_constraint(t=0, derivative_order=2, lb=[0, 0])\n",
    "zpp.add_constraint(t=1, derivative_order=0, lb=[2, 1.5])\n",
    "zpp.add_constraint(t=2, derivative_order=0, lb=[4, 1])\n",
    "zpp.add_constraint(t=tf, derivative_order=0, lb=[6, 0])\n",
    "zpp.add_constraint(t=tf, derivative_order=1, lb=[0, 0])\n",
    "zpp.add_constraint(t=tf, derivative_order=2, lb=[0, 0])\n",
    "zpp.generate()\n",
    "\n",
    "if False:  # Useful for debugging\n",
    "    t = np.linspace(0, tf, 100)\n",
    "    z = np.zeros((2, len(t)))\n",
    "    knots = np.zeros((2, len(zpp.sample_times)))\n",
    "    fig, ax = plt.subplots(zpp.degree + 1, 1)\n",
    "    for deg in range(zpp.degree + 1):\n",
    "        for i in range(len(t)):\n",
    "            z[:, i] = zpp.eval(t[i], deg)\n",
    "        for i in range(len(zpp.sample_times)):\n",
    "            knots[:, i] = zpp.eval(zpp.sample_times[i], deg)\n",
    "        ax[deg].plot(t, z.transpose())\n",
    "        ax[deg].plot(zpp.sample_times, knots.transpose(), \".\")\n",
    "        ax[deg].set_xlabel(\"t (sec)\")\n",
    "        ax[deg].set_ylabel(\"z deriv \" + str(deg))\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "t = np.linspace(0, tf, 100)\n",
    "z = np.zeros((2, len(t)))\n",
    "\n",
    "for i in range(len(t)):\n",
    "    z[:, i] = zpp.eval(t[i])\n",
    "ax.plot(z[0, :], z[1, :])\n",
    "\n",
    "for t in np.linspace(0, tf, 7):\n",
    "    x = zpp.eval(t)\n",
    "    xddot = zpp.eval(t, 2)\n",
    "    theta = np.arctan2(-xddot[0], (xddot[1] + 9.81))\n",
    "    v = Quadrotor2DVisualizer(ax=ax)\n",
    "    context = v.CreateDefaultContext()\n",
    "    context.FixInputPort(0, [x[0], x[1], theta, 0, 0, 0])\n",
    "    v.draw(context)\n",
    "\n",
    "# Draw the (imaginary) obstacles\n",
    "ax.fill(2 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.25 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(2 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.75 + 1.25 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(4 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        .75 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.fill(4 + np.array([-.1, -.1, .1, .1, -.1]),\n",
    "        1.25 + 1.75 * np.array([0, 1, 1, 0, 0]),\n",
    "        facecolor=\"darkred\",\n",
    "        edgecolor=\"k\")\n",
    "ax.set_xlim([-1, 7])\n",
    "ax.set_ylim([-.25, 3])\n",
    "ax.set_title(\"\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
